/*** Created Time: 2022年5月16日 18:56:46 ***/
/*** Author: Eric                        ***/

// 多任务调度实现

#include <onix/task.h>
#include <onix/printk.h>
#include <onix/debug.h>
#include <onix/memory.h>
#include <onix/assert.h>
#include <onix/interrupt.h>
#include <onix/string.h>
#include <onix/bitmap.h>
#include <onix/syscall.h>
#include <onix/list.h>
#include <onix/global.h>
#include <onix/arena.h>

#define     LOGK(fmt, args...) DEBUGK(fmt, ##args)

#define     NR_TASKS    (64)  // 任务数量

extern u32 volatile jiffies;            // 导入全局时间片数量，定义于clock.c
extern u32 jiffy;                       // 导入一个时间片的毫秒值
extern bitmap_t kernel_map;             // 获得内核位图
extern tss_t tss;                       // 获得任务状态段

extern void task_switch(task_t *next);  // 任务切换函数，在schedule.asm中实现

static task_t *task_table[NR_TASKS];    // 任务表
static list_t block_list;               // 任务默认阻塞链表
static list_t sleep_list;               // 任务睡眠链表，一定有序

static task_t *idle_task;               // 空闲任务

// 从 task_table 里获得一个空闲的任务
static task_t *get_free_task(void)
{
    for(size_t i = 0; i < NR_TASKS; i++) {
        if(task_table[i] == NULL) {
            // 一个空闲任务可以分配一页页表
            task_t *task = (task_t *)alloc_kpage(1); // todo free_kpage
            memset(task, 0, PAGE_SIZE);
            task->pid = i;
            task_table[i] = task;
            return task;
        }
    }
    panic("No more tasks");
}

// 获取进程 id
pid_t sys_getpid(void)
{
    task_t *task = running_task();
    return task->pid;
}

// 获取父进程 id
pid_t sys_getppid(void)
{
    task_t *task = running_task();
    return task->ppid;
}

// 从任务数组中查找某种状态的任务，自己除外
static task_t *task_search(task_state_t state)
{
    assert(!get_interrupt_state());     // 保证是在不可中断的情况下，即确保原子操作
    task_t *task = NULL;
    task_t *current = running_task();

    for(size_t i = 0; i < NR_TASKS; i++) {
        task_t *ptr = task_table[i];
        if(ptr == NULL) {
            continue;
        }
        if(ptr->state != state) {   // 查找某种状态的任务
            continue;
        }
        if(ptr == current) {    // 自己除外
            continue;
        }
        if(task == NULL || task->ticks < ptr->ticks || ptr->jiffies < task->jiffies) {
            task = ptr;
        }
    }

    // 如果找不到处于就绪状态的任务，则将task修改为空闲任务并返回
    if(task == NULL && state == TASK_READY) {
        task = idle_task;
    }

    return task;
}

void task_yield(void)
{
    schedule();
}

// 任务阻塞
// @task : 待阻塞任务
// @blist : 任务阻塞链表
// @state : 阻塞状态
void task_block(task_t *task, list_t *blist, task_state_t state)
{
    assert(!get_interrupt_state());     // 确保处于不可中断的条件下
    assert(task->node.next == NULL);    // 确保当前task没有加入任何一个链表中
    assert(task->node.prev == NULL);

    if (blist == NULL) {    // 如果blist为空，则赋值为默认的任务阻塞链表
        blist = &block_list;
    }

    list_push(blist, &task->node);

    // 如果task要阻塞，则不能将task改为READY或者RUNNING状态
    assert(state != TASK_READY && state != TASK_RUNNING);
    task->state = state;

    task_t *current = running_task();
    if (current == task) {
        // 如果当前运行任务就是task，则调度执行其他任务
        schedule();
    }
}

// 解除任务阻塞
void task_unblock(task_t *task)
{
    assert(!get_interrupt_state());     // 确保处于不可中断的条件下

    list_remove(&task->node);           // 将task从阻塞链表中移除

    assert(task->node.next == NULL);
    assert(task->node.prev == NULL);

    task->state = TASK_READY;           // 返回就绪态
}

// 任务睡眠函数
void task_sleep(u32 ms)
{
    assert(!get_interrupt_state()); // 不可中断

    u32 ticks = ms / jiffy;        // 需要睡眠的时间片
    ticks = ticks > 0 ? ticks : 1; // 至少休眠一个时间片

    // 记录目标全局时间片，在那个时刻需要唤醒任务
    task_t *current = running_task();
    current->ticks = jiffies + ticks;   // 当全局时间片到达jiffies + ticks时就可以让任务就绪

    // 从睡眠链表找到第一个比当前任务唤醒时间点更晚的任务，进行插入排序，插入到该任务前面，确保睡眠列表一定有序
    list_t *list = &sleep_list;
    list_node_t *anchor = &list->tail;
    // 插入排序实现
    for (list_node_t *ptr = list->head.next; ptr != &list->tail; ptr = ptr->next) {
        // 通过ptr(list_node_t *)类型来求task_t结构体的起始地址
        task_t *task = element_entry(task_t, node, ptr);

        if (task->ticks > current->ticks) {
            anchor = ptr;
            break;
        }
    }

    // 确保当前运行的任务没有被加入到任何一个链表中
    assert(current->node.next == NULL);
    assert(current->node.prev == NULL);

    // 插入链表
    list_insert_before(anchor, &current->node);

    // 阻塞状态是睡眠
    current->state = TASK_SLEEPING;

    // 调度执行其他任务
    schedule();
}

// 任务唤醒函数
void task_wakeup(void)
{
    assert(!get_interrupt_state()); // 不可中断

    // 从睡眠链表中找到 ticks 小于等于 jiffies 的任务，恢复执行
    list_t *list = &sleep_list;
    for (list_node_t *ptr = list->head.next; ptr != &list->tail; /* 空 */) {
        // 通过ptr(list_node_t *)类型来求task_t结构体的起始地址
        task_t *task = element_entry(task_t, node, ptr);
        if (task->ticks > jiffies) {
            break;
        }

        // unblock 会将指针清空
        ptr = ptr->next;

        task->ticks = 0;
        task_unblock(task);
    }
}

// 激活任务
void task_activate(task_t *task)
{
    assert(task->magic == ONIX_MAGIC);

    if (task->pde != get_cr3()) {
        // 将页表切换到新的任务
        set_cr3(task->pde);
        // BMB;
    }

    // 如果 task 的用户不是内核
    if (task->uid != KERNEL_USER)
    {
        tss.esp0 = (u32)task + PAGE_SIZE;
    }
}

// 当前运行的任务/进程，%esp存放当前进程的栈顶
task_t *running_task(void)
{
    __asm__ __volatile__ (
        "movl %esp, %eax\n"
        "andl $0xfffff000, %eax\n"
    );
}

// 任务调度函数
void schedule(void)
{
    assert(!get_interrupt_state());         // 不可中断

    task_t *current = running_task();
    task_t *next = task_search(TASK_READY); // 查找处于就绪状态的任务

    assert(next != NULL);
    assert(next->magic == ONIX_MAGIC);

    if(current->state == TASK_RUNNING) {
        current->state = TASK_READY;    // 将当前处于运行状态的任务切换为就绪态
    }

    if(!current->ticks) {
        current->ticks = current->priority; // 如果当前任务时间片为空，则赋为当前优先级
    }

    next->state = TASK_RUNNING;         // 从就绪态切换到运行态
    if(next == current) {
        return;         // 如果下一个任务和当前任务相同，则不执行任务切换操作
    }

    task_activate(next);
    task_switch(next);  // 执行真正的任务切换
}

// 创建任务函数
// @target : 待创建的任务
// @name : 任务名字
// @priority : 任务优先级
// @uid : 任务的ID，内核任务 / 普通任务
static task_t *task_create(target_t target, const char *name, u32 priority, u32 uid)
{
    task_t *task = get_free_task();

    u32 stack = (u32)task + PAGE_SIZE;  // 内核栈栈顶位于页表的最底部

    stack -= sizeof(task_frame_t);
    task_frame_t *frame = (task_frame_t *)stack;
    frame->ebx = 0x11111111;
    frame->esi = 0x22222222;
    frame->edi = 0x33333333;
    frame->ebp = 0x44444444;
    frame->eip = (void *)target;

    strcpy((char *)task->name, name);

    task->stack = (u32 *)stack;     // 设置任务栈顶地址
    task->priority = priority;      // 设置任务优先级
    task->ticks = task->priority;   // 设置任务时间片，初始时就等于优先级
    task->jiffies = 0;              // 上次执行时时间片为0
    task->state = TASK_READY;       // 就绪态
    task->uid = uid;                // 用户ID
    task->vmap = &kernel_map;       // 内核虚拟内存的位图
    task->pde = KERNEL_PAGE_DIR;    // 内核页目录地址 page directory entry
    task->brk = KERNEL_MEMORY_SIZE; // 进程堆地址最高地址设置为 8M 
    task->magic = ONIX_MAGIC;       // 防止栈溢出，类似于金丝雀

    return task;
}

// 调用该函数的地方不能有任何局部变量
// 调用前栈顶需要准备足够的空间
void task_to_user_mode(target_t target)
{
    task_t *task = running_task();

    // 创建用户进程虚拟内存位图
    task->vmap = kmalloc(sizeof(bitmap_t)); 
    void *buf = (void *)alloc_kpage(1); 
    bitmap_init(task->vmap, buf, PAGE_SIZE, KERNEL_MEMORY_SIZE / PAGE_SIZE);

    // 创建用户进程页表
    task->pde = (u32)copy_pde();
    set_cr3(task->pde);

    u32 addr = (u32)task + PAGE_SIZE;

    // 先求出栈顶，然后求出减去一个中断帧大小后的新栈顶
    addr -= sizeof(intr_frame_t);
    intr_frame_t *iframe = (intr_frame_t *)(addr);

    iframe->vector = 0x20;
    iframe->edi = 1;
    iframe->esi = 2;
    iframe->ebp = 3;
    iframe->esp_dummy = 4;
    iframe->ebx = 5;
    iframe->edx = 6;
    iframe->ecx = 7;
    iframe->eax = 8;

    iframe->gs = 0;
    iframe->ds = USER_DATA_SELECTOR;
    iframe->es = USER_DATA_SELECTOR;
    iframe->fs = USER_DATA_SELECTOR;
    iframe->ss = USER_DATA_SELECTOR;
    iframe->cs = USER_CODE_SELECTOR;

    iframe->error = ONIX_MAGIC;

    iframe->eip = (u32)target;
    iframe->eflags = (0 << 12 | 0b10 | 1 << 9);
    iframe->esp = USER_STACK_TOP;   // 用户态栈顶

    asm volatile(
        "movl %0, %%esp\n"
        "jmp interrupt_exit\n" ::"m"(iframe));
}

extern void interrupt_exit();

// 构造 fork 系统调用后的子进程栈
static void task_build_stack(task_t *task)
{
    u32 addr = (u32)task + PAGE_SIZE;
    addr -= sizeof(intr_frame_t);
    intr_frame_t *iframe = (intr_frame_t *)addr;
    iframe->eax = 0;    // 子进程返回 0 

    addr -= sizeof(task_frame_t);
    task_frame_t *frame = (task_frame_t *)addr;

    frame->ebp = 0xaa55aa55;
    frame->ebx = 0xaa55aa55;
    frame->edi = 0xaa55aa55;
    frame->esi = 0xaa55aa55;

    frame->eip = interrupt_exit;

    task->stack = (u32 *)frame;
}

// fork 系统调用
pid_t task_fork(void)
{
    // LOGK("fork is called\n");
    task_t *task = running_task();

    // 当前进程没有阻塞，且正在执行
    assert(task->node.next == NULL && task->node.prev == NULL && task->state == TASK_RUNNING);

    // 拷贝内核栈和 PCB
    task_t *child = get_free_task();
    pid_t pid = child->pid;
    memcpy(child, task, PAGE_SIZE);

    child->pid = pid;
    child->ppid = task->pid;
    child->ticks = child->priority;
    child->state = TASK_READY;

    // 拷贝用户进程虚拟内存位图
    child->vmap = kmalloc(sizeof(bitmap_t));
    memcpy(child->vmap, task->vmap, sizeof(bitmap_t));

    // 拷贝虚拟位图缓存
    void *buf = (void *)alloc_kpage(1);
    memcpy(buf, task->vmap->bits, PAGE_SIZE);
    child->vmap->bits = buf;

    // 拷贝页目录
    child->pde = (u32)copy_pde();

    // 构造 child 内核栈
    task_build_stack(child); // ROP
    // schedule();

    return child->pid;  // 父进程返回子进程的 pid
}

// exit 系统调用
void task_exit(int status)
{
    task_t *task = running_task();

    // 当前进程没有阻塞，且正在执行
    assert(task->node.next == NULL && task->node.prev == NULL && task->state == TASK_RUNNING);

    task->state = TASK_DIED;
    task->status = status;

    // 释放页目录
    free_pde();

    // 释放虚拟内存位图
    free_kpage((u32)task->vmap->bits, 1);
    kfree(task->vmap);

    // 将子进程的父进程赋值为自己的父进程
    for (size_t i = 2; i < NR_TASKS; i++) {
        task_t *child = task_table[i];
        if (!child)
            continue;
        if (child->ppid != task->pid)
            continue;
        child->ppid = task->ppid;
    }
    LOGK("task 0x%p exit....\n", task);

    // waitpid 相关
    task_t *parent = task_table[task->ppid];
    if (parent->state == TASK_WAITING &&
        (parent->waitpid == -1 || parent->waitpid == task->pid))
    {
        task_unblock(parent);
    }

    schedule();
}

pid_t task_waitpid(pid_t pid, int32 *status)
{
    task_t *task = running_task();
    task_t *child = NULL;

    while (true)
    {
        // 检测是否有子进程需要等待
        bool has_child = false;
        for (size_t i = 2; i < NR_TASKS; i++) {
            task_t *ptr = task_table[i];
            if (!ptr)
                continue;

            if (ptr->ppid != task->pid)
                continue;
            if (pid != ptr->pid && pid != -1)
                continue;

            if (ptr->state == TASK_DIED) {
                child = ptr;
                task_table[i] = NULL;
                goto rollback;
            }

            has_child = true;
        }
        if (has_child) {
            task->waitpid = pid;
            task_block(task, NULL, TASK_WAITING);
            continue;
        }
        break;
    }

    // 没找到符合条件的子进程
    return -1;

rollback:
    *status = child->status;
    u32 ret = child->pid;
    free_kpage((u32)child, 1);
    return ret;
}

static void task_setup(void)
{
    task_t *task = running_task();
    task->magic = ONIX_MAGIC;
    task->ticks = 1;

    memset(task_table, 0, sizeof(task_table));
}

extern void idle_thread(void);
extern void init_thread(void);
extern void test_thread(void);

void task_init(void)
{
    list_init(&block_list);     // 初始化阻塞链表
    list_init(&sleep_list);     // 初始化睡眠链表

    task_setup();

    idle_task = task_create(idle_thread, "idle", 1, KERNEL_USER);   // 创建空闲任务
    task_create(init_thread, "init", 5, NORMAL_USER);
    task_create(test_thread, "test", 5, NORMAL_USER);
}